/*
 * Copyright 2025 author or author.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.jcwang.commons.tools.util.encryption.rsa;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import org.apache.commons.codec.binary.Base64;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.lang.NonNull;
import org.springframework.util.Assert;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.time.Duration;
import java.util.Arrays;

/**
 * RSA加解密工具类
 *
 * @author 王金城
 **/
public class RsaUtils {

    /**
     * Key大小
     **/
    private static final int KEY_SIZE = 2048;
    /**
     * 最大加密字节数。<br>
     * 不管明文长度是多少，RSA 生成的密文长度总是固定的。
     * 但是明文长度不能超过密钥长度。比如 Java 默认的 RSA 加密实现不允许明文长度超过密钥长度减去 11(单位是字节，也就是 byte)。也就是说，
     * 如果我们定义的密钥(我们可以通过 java.security.KeyPairGenerator.initialize(int keySize) 来定义密钥长度)长度为 2048(单位是位，也就是 bit)，
     * 生成的密钥长度就是 2048位 / 8位/字节 = 256字节，那么我们需要加密的明文长度不能超过 256字节 - 11 字节 = 245字节。也就是说，
     * 我们最大能将 245 字节长度的明文进行加密，否则会出问题(抛诸如 javax.crypto.IllegalBlockSizeException: Data must not be longer than 245 bytes 的异常)。
     * 而 BC 提供的加密算法能够支持到的 RSA 明文长度最长为密钥长度。
     **/
    private static final int MAX_ENCRYPT_BLOCK = KEY_SIZE / 8 - 11;
    /**
     * 最大解密字节数。与加密不同。最大长度为 KEY_SIZE 的 1/8
     **/
    private static final int MAX_DECRYPT_BLOCK = KEY_SIZE / 8;
    /**
     * 算法名称
     **/
    private static final String ALGORITHM_NAME = "RSA";
    /**
     * 在Cipher中的算法名称
     **/
    private static final String CIPHER_ALGORITHM_NAME = "RSA/ECB/PKCS1Padding";
    /**
     * 签名算法
     **/
    private static final String SIGN_ALGORITHM = "SHA256withRSA";
    /**
     * 公钥缓存
     **/
    private static final Cache<String, PublicKey> PUBLIC_KEY_CACHE;
    /**
     * 私钥缓存
     **/
    private static final Cache<String, PrivateKey> PRIVATE_KEY_CACHE;
    private static final Logger logger = LoggerFactory.getLogger(RsaUtils.class);

    static {
        int i = Security.addProvider(new BouncyCastleProvider());
        logger.info("RSA工具类添加提供者 --- [Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider()] - 提供程序的首选项位置 {}", i);
        PUBLIC_KEY_CACHE = Caffeine.newBuilder()
                .expireAfterAccess(Duration.ofHours(8))
                .maximumSize(128)
                .build();
        PRIVATE_KEY_CACHE = Caffeine.newBuilder()
                .expireAfterAccess(Duration.ofHours(8))
                .maximumSize(128)
                .build();
        logger.info("RSA 公钥、私钥缓存初始化成功 - [maxSize: 128, expireAfterAccess: 8h]；已完成对RSA工具类的初始化...");
    }

    /**
     * 获取密钥对
     *
     * @return java.security.KeyPair
     * @author 王金城
     */
    public static KeyPair getKeyPair() {
        try {
            KeyPairGenerator generator = KeyPairGenerator.getInstance(ALGORITHM_NAME);
            generator.initialize(KEY_SIZE);
            return generator.generateKeyPair();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("没有找到算法", e);
        }
    }

    /**
     * 获取string类型的密钥对
     *
     * @return com.jcwang.commons.util.encryption.rsa.RsaStringKeyPair
     * @author 王金城
     */
    public static RsaStringKeyPair getKeyPairString() {
        KeyPair keyPair = getKeyPair();
        byte[] privateKeyBytes = keyPair.getPrivate().getEncoded();
        String privateKey = new String(Base64.encodeBase64(privateKeyBytes), StandardCharsets.UTF_8);
        byte[] publicKeyBytes = keyPair.getPublic().getEncoded();
        String publicKey = new String(Base64.encodeBase64(publicKeyBytes), StandardCharsets.UTF_8);
        return new RsaStringKeyPair(privateKey, publicKey);
    }

    /**
     * 获取公钥
     *
     * @param publicKey base64加密的公钥字符串
     * @return java.security.PublicKey
     * @author 王金城
     */
    public static PublicKey getPublicKey(@NonNull String publicKey) {
        Assert.hasText(publicKey, "公钥不能为空");
        return PUBLIC_KEY_CACHE.get(publicKey, (k) -> {
            try {
                byte[] decodedKey = Base64.decodeBase64(k);
                X509EncodedKeySpec keySpec = new X509EncodedKeySpec(decodedKey);
                KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_NAME);
                return keyFactory.generatePublic(keySpec);
            } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
                throw new RuntimeException("获取公钥异常", e);
            }
        });
    }

    /**
     * RSA加密，默认非 URL 安全的 Base64 字符串
     *
     * @param source 待加密数据
     * @param key    密钥
     * @return 非 URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String encrypt(@NonNull String source, @NonNull Key key) {
        return encrypt(source, key, false);
    }

    /**
     * RSA加密，返回 URL 安全的 Base64 字符串
     *
     * @param source 待加密数据
     * @param key    密钥
     * @return URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String encryptURLSafe(@NonNull String source, @NonNull Key key) {
        return encrypt(source, key, true);
    }

    /**
     * RSA加密
     *
     * @param source  待加密数据
     * @param key     密钥
     * @param urlSafe 是否需要生成 URL 安全的 Base64 字符串
     * @return java.lang.String
     * @author 王金城
     */
    public static String encrypt(@NonNull String source, @NonNull Key key, boolean urlSafe) {
        Assert.notNull(source, "加密数据不能为 null");
        Assert.notNull(key, "密钥不能为 null");
        try {
            Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM_NAME);
            cipher.init(Cipher.ENCRYPT_MODE, key);
            byte[] sourceBytes = source.getBytes();
            byte[] bytes = doCipherFinal(sourceBytes, cipher, MAX_ENCRYPT_BLOCK);
            if (urlSafe) {
                return new String(Base64.encodeBase64URLSafe(bytes), StandardCharsets.UTF_8);
            }
            return new String(Base64.encodeBase64(bytes), StandardCharsets.UTF_8);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException |
                 BadPaddingException e) {
            throw new RuntimeException("RSA加密异常", e);
        }
    }


    /**
     * RSA加密
     *
     * @param source    待加密数据
     * @param publicKey 公钥
     * @return 非 URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String encrypt(@NonNull String source, @NonNull String publicKey) {
        Assert.notNull(source, "加密数据不能为 null");
        PublicKey key = getPublicKey(publicKey);
        return encrypt(source, key);
    }

    /**
     * RSA加密
     *
     * @param source    待加密数据
     * @param publicKey 公钥
     * @return URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String encryptURLSafe(@NonNull String source, @NonNull String publicKey) {
        Assert.notNull(source, "加密数据不能为 null");
        PublicKey key = getPublicKey(publicKey);
        return encryptURLSafe(source, key);
    }

    /**
     * RSA加密（反向使用，使用私钥加密，公钥解密）
     *
     * @param source     待加密数据
     * @param privateKey 私钥
     * @return java.lang.String
     * @author 王金城
     */
    public static String encryptFromPrivateKey(@NonNull String source, @NonNull String privateKey) {
        Assert.notNull(source, "加密数据不能为 null");
        PrivateKey key = getPrivateKey(privateKey);
        return encrypt(source, key);
    }

    /**
     * RSA加密（反向使用，使用私钥加密，公钥解密）
     *
     * @param source     待加密数据
     * @param privateKey 私钥
     * @return URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String encryptURLSafeFromPrivateKey(@NonNull String source, @NonNull String privateKey) {
        Assert.notNull(source, "加密数据不能为 null");
        PrivateKey key = getPrivateKey(privateKey);
        return encryptURLSafe(source, key);
    }

    /**
     * 获取私钥
     *
     * @param privateKey base64加密的私钥字符串
     * @return java.security.PrivateKey
     * @author 王金城
     */
    public static PrivateKey getPrivateKey(@NonNull String privateKey) {
        Assert.hasText(privateKey, "私钥不能为空");
        return PRIVATE_KEY_CACHE.get(privateKey, (k) -> {
            try {
                byte[] decodedKey = Base64.decodeBase64(k);
                PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(decodedKey);
                KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_NAME);
                return keyFactory.generatePrivate(keySpec);
            } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
                throw new RuntimeException("获取私钥异常", e);
            }
        });
    }

    /**
     * RSA解密
     *
     * @param source 待解密数据
     * @param key    密钥
     * @return java.lang.String
     * @author 王金城
     */
    public static String decrypt(@NonNull String source, @NonNull Key key) {
        Assert.hasText(source, "解密数据不能为空");
        Assert.notNull(key, "密钥不能为 null");
        try {
            Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM_NAME);
            cipher.init(Cipher.DECRYPT_MODE, key);
            byte[] sourceBytes = Base64.decodeBase64(source);
            byte[] bytes = doCipherFinal(sourceBytes, cipher, MAX_DECRYPT_BLOCK);
            return new String(bytes, StandardCharsets.UTF_8);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException |
                 BadPaddingException e) {
            throw new RuntimeException("RSA解密异常", e);
        }
    }

    /**
     * RSA解密
     *
     * @param source     待解密数据
     * @param privateKey 私钥
     * @return java.lang.String
     * @author 王金城
     */
    public static String decrypt(@NonNull String source, @NonNull String privateKey) {
        Assert.hasText(source, "解密数据不能为空");
        PrivateKey key = getPrivateKey(privateKey);
        return decrypt(source, key);
    }

    /**
     * RSA解密（反向使用，使用私钥加密，公钥解密）
     *
     * @param source    待解密数据
     * @param publicKey 公钥
     * @return java.lang.String
     * @author 王金城
     */
    public static String decryptFromPublicKey(@NonNull String source, @NonNull String publicKey) {
        Assert.hasText(source, "解密数据不能为空");
        PublicKey key = getPublicKey(publicKey);
        return decrypt(source, key);
    }

    /**
     * RSA 签名
     *
     * @param source     待签名数据
     * @param privateKey 私钥
     * @return 非 URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String sign(@NonNull String source, @NonNull PrivateKey privateKey) {
        return sign(source, privateKey, false);
    }

    /**
     * RSA 签名
     *
     * @param source     待签名数据
     * @param privateKey 私钥
     * @return URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String signURLSafe(@NonNull String source, @NonNull PrivateKey privateKey) {
        return sign(source, privateKey, true);
    }

    /**
     * RSA 签名
     *
     * @param source     待签名数据
     * @param privateKey 私钥
     * @param urlSafe    是否使用 URL 安全的 Base64 编码
     * @return java.lang.String
     * @author 王金城
     */
    public static String sign(@NonNull String source, @NonNull PrivateKey privateKey, boolean urlSafe) {
        Assert.hasText(source, "待签名数据不能为空");
        Assert.notNull(privateKey, "私钥不能为 null");
        try {
            byte[] sourceBytes = source.getBytes();
            Signature signature = Signature.getInstance(SIGN_ALGORITHM);
            signature.initSign(privateKey);
            signature.update(sourceBytes);
            if (urlSafe) {
                return new String(Base64.encodeBase64URLSafe(signature.sign()), StandardCharsets.UTF_8);
            }
            return new String(Base64.encodeBase64(signature.sign()), StandardCharsets.UTF_8);
        } catch (NoSuchAlgorithmException | InvalidKeyException | SignatureException e) {
            throw new RuntimeException("RSA 签名异常", e);
        }
    }

    /**
     * RSA 签名
     *
     * @param source     待签名数据
     * @param privateKey 私钥
     * @return 非 URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String sign(@NonNull String source, @NonNull String privateKey) {
        Assert.hasText(privateKey, "私钥不能为空");
        PrivateKey key = getPrivateKey(privateKey);
        return sign(source, key);
    }

    /**
     * RSA 签名
     *
     * @param source     待签名数据
     * @param privateKey 私钥
     * @return URL 安全的 Base64 字符串
     * @author 王金城
     */
    public static String signURLSafe(@NonNull String source, @NonNull String privateKey) {
        Assert.hasText(privateKey, "私钥不能为空");
        PrivateKey key = getPrivateKey(privateKey);
        return signURLSafe(source, key);
    }

    /**
     * RSA 验证签名
     *
     * @param source    待验签数据
     * @param sign      签名
     * @param publicKey 公钥
     * @return boolean
     * @author 王金城
     */
    public static boolean verify(@NonNull String source, @NonNull String sign, @NonNull PublicKey publicKey) {
        Assert.notNull(publicKey, "公钥不能为 null");
        Assert.hasText(source, "验签数据不能为空");
        Assert.hasText(sign, "签名不能为空");

        try {
            byte[] sourceBytes = source.getBytes();
            byte[] signBytes = Base64.decodeBase64(sign);
            Signature signature = Signature.getInstance(SIGN_ALGORITHM);
            signature.initVerify(publicKey);
            signature.update(sourceBytes);
            return signature.verify(signBytes);
        } catch (NoSuchAlgorithmException | InvalidKeyException | SignatureException e) {
            throw new RuntimeException("RSA 验证签名异常", e);
        }
    }

    /**
     * RSA 验证签名
     *
     * @param source    待验签数据
     * @param sign      签名
     * @param publicKey 公钥
     * @return boolean
     * @author 王金城
     */
    public static boolean verify(@NonNull String source, @NonNull String sign, @NonNull String publicKey) {
        Assert.hasText(publicKey, "公钥不能为空");
        PublicKey key = getPublicKey(publicKey);
        return verify(source, sign, key);
    }

    /**
     * 加解密数据
     *
     * @param sourceBytes 源数据byte数组
     * @param cipher      加解密类实例
     * @param maxBlock    分段加解密时，最大块大小
     * @return byte[]
     * @author 王金城
     */
    private static byte[] doCipherFinal(byte[] sourceBytes, Cipher cipher, int maxBlock) throws IllegalBlockSizeException, BadPaddingException {
        int inputLength = sourceBytes.length;
        // 如果长度未超过最大长度，则直接返回结果，不进行分段处理
        if (inputLength <= maxBlock) {
            return cipher.doFinal(sourceBytes);
        }
        // 开始处的输入偏移
        int inputOffset = 0;
        // 结果数组
        byte[] resultBytes = {};
        // 分段加密或解密的数组
        byte[] cache;
        while (inputLength - inputOffset > 0) {
            if (inputLength - inputOffset > maxBlock) {
                cache = cipher.doFinal(sourceBytes, inputOffset, maxBlock);
                inputOffset += maxBlock;
            } else {
                cache = cipher.doFinal(sourceBytes, inputOffset, inputLength - inputOffset);
                inputOffset = inputLength;
            }
            // 复制结果数组，并扩展数组长度
            resultBytes = Arrays.copyOf(resultBytes, resultBytes.length + cache.length);
            // 将cache数组加入结果数组中
            System.arraycopy(cache, 0, resultBytes, resultBytes.length - cache.length, cache.length);
        }
        return resultBytes;
    }
}
